////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                        //
//---------------------------- Ebrahim Foulaadvand, 05 Aug 2013 -----------------                         //
//                                                                                                        //
//  The routine "SchroFTCS" solves the 1D Schrodinger wave equation using forward time centered space     //
//  scheme. Periodic boundary condition is used. The system length L and we take hbar=m=1.                //
//  Initial wave shape is cosine-modulated Gaussian pulse.                                                //                                                                                                                                                              //
//                                                                                                        //                                                                                                        //
////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include <iostream>
#include <fstream>
#include <conio>
#include <math.h>
#include <numeric>
#include <values>
#include <vector>
#include <list>
#include <algorithm>
#include <numeric>

using namespace std;

main()
{

double hbar=1,m=1,tau=0.005,h,L=40,sigma0=0.05,k0=0.2,x,x0=0,p0=hbar*(k0/m),
       taw=2*m*sigma0*sigma0/hbar;

int N=200,i,n,T=1000;


ofstream file0 ("initial profile n=0.plt");     //output file for the probability profile at timestep n=0.
ofstream file1 ("probability FTCS n=10.plt");    //output file for the probability profile at timestep n=1.
ofstream file2 ("probability FTCS n=20.plt");    //output file for the probability profile at timestep n=2.
ofstream file3 ("probability anal n=10.plt");    //output file for theabalytic at timestep n=1.
ofstream file4 ("probability anal n=20.plt");    //output file for theabalytic at timestep n=1.


vector<double> PsiR(N+1,0),PsiRnew(N+1,0),PsiI(N+1,0),PsiInew(N+1,0),Probanal(N+1,0),V(N+1,0);

 // Arrays "PsiR" and "PsiI" store the current values of the real and imaginary parts of the wavefunction.
 // Arrays "PsiRnew", "PsiInew" store the updated values of the real and imaginary parts of the wavefunction.


 h=L/N; cout<<"h= "<<h<<endl; //getch();

 cout<<"tau= "<<tau<<endl; //getch();

 cout<<"taw= "<<taw<<endl; cout<<"po0= "<<p0<<endl;


for(i=0;i<=N;i++){ x=-0.5*L+i*h; V[i]=0; // we consider free particle for which V(x)=0

PsiR[i]=( 1/sqrt(sigma0*sqrt(M_PI)) )*cos(k0*x)*exp(-0.5*(x/sigma0)*(x/sigma0));

PsiI[i]=( 1/sqrt(sigma0*sqrt(M_PI)) )*sin(k0*x)*exp(-0.5*(x/sigma0)*(x/sigma0));

file0<<(-0.5*L+i*h)<<"  "<<sqrt(PsiR[i]*PsiR[i]+PsiI[i]*PsiI[i])<<endl; }

//storing the initial condition data in the data file.


for (n=1;n<=T;n++){ //cout<<"n= "<<n<<endl; //getch();


// -------------------       FTCS method -------------------------------------


PsiRnew[0]=PsiR[0]-(tau/(2*h*h))*(PsiI[1]+PsiI[N]-2*PsiI[0]) + tau*V[0]*PsiR[0];

PsiInew[0]=PsiI[0]+(tau/(2*h*h))*(PsiR[1]+PsiR[N]-2*PsiR[0]) - tau*V[0]*PsiI[0]; // periodic condition has been used: u[-1] is replaced with u[N]


for(i=1;i<N;i++){

PsiRnew[i]=PsiR[i]-(tau/(2*h*h))*(PsiI[i+1]+PsiI[i-1]-2*PsiI[i]) + tau*V[i]*PsiR[i];

PsiInew[i]=PsiI[i]+(tau/(2*h*h))*(PsiR[i+1]+PsiR[i-1]-2*PsiR[i]) - tau*V[i]*PsiI[i];

 }


PsiRnew[N]=PsiR[N]-(tau/(2*h*h))*(PsiI[0]+PsiI[N-1]-2*PsiI[N]) + tau*V[N]*PsiR[N];

PsiInew[N]=PsiI[N]+(tau/(2*h*h))*(PsiR[0]+PsiR[N-1]-2*PsiR[N]) - tau*V[N]*PsiI[N];// periodic condition has been used: u[N+1] is replaced with u[0]


if (n==10) {


for(i=0;i<=N;i++){

file1<<(-0.5*L+i*h)<<"  "<<sqrt(PsiRnew[i]*PsiRnew[i]+PsiInew[i]*PsiInew[i])<<endl;

x=-0.5*L+i*h;

Probanal[i]=(1/(sigma0*sqrt(2*M_PI*(1+(n*tau/taw)*(n*tau/taw)))))*exp(-(1/(2*sigma0*sigma0*sqrt((1+(n*tau/taw)*(n*tau/taw)))))*(x-(p0/m)*n*tau)*(x-(p0/m)*n*tau));

cout<<x-(p0/m)*n*tau<<endl;getch();

//cout<<"Probanal["<<i<<"] "<<Probanal[i]<<endl;

file3<<(-0.5*L+i*h)<<"  "<<Probanal[i]<<endl;

} //storing in the data file.
}


if (n==20) {


for(i=0;i<=N;i++){

file2<<(-0.5*L+i*h)<<"  "<<sqrt(PsiRnew[i]*PsiRnew[i]+PsiInew[i]*PsiInew[i])<<endl;

x=-0.5*L+i*h;

Probanal[i]=(1/(sigma0*sqrt(2*M_PI*(1+(n*tau/taw)*(n*tau/taw)))))*exp(-(1/(2*sigma0*sigma0*sqrt((1+(n*tau/taw)*(n*tau/taw)))))*(x-(p0/m)*n*tau)*(x-(p0/m)*n*tau));

file4<<(-0.5*L+i*h)<<"  "<<Probanal[i]<<endl;


} //storing in the data file.
}



for(i=0;i<=N;i++){ PsiR[i]=PsiRnew[i]; PsiI[i]=PsiInew[i]; }


}

cout<<"end." <<endl; getch();

}